
#include "vector/cross_entropy_loss_test.h"

#include "log.h"
#include "vector/normal.h"
#include "vector/cross_entropy_loss.h"
#include "vector/tensor.h"
#include "tools/measure_execution_time.h"

static const std::vector<std::vector<double>> input{
{0.00996962, 0.064692, 0.00775108, 0.221152, 0.419546, 0.00238489, 0.0262764, 0.200586, 0.0406673, 0.00697486}, 
{0.0148357, 0.0284819, 0.012694, 0.164776, 0.589914, 0.0101577, 0.0459329, 0.0346252, 0.0841187, 0.0144633}, 
{0.0204002, 0.0370894, 0.0392826, 0.26788, 0.114675, 0.00216199, 0.0116182, 0.0927034, 0.397728, 0.0164611}, 
{0.0319662, 0.0329964, 0.00971772, 0.276478, 0.428359, 0.0113864, 0.0445876, 0.0278179, 0.118963, 0.0177281}};

static const std::vector<std::vector<double>> labels{
{1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}, 
{0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}, 
{0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}, 
{0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}};

static const double expect_loss{12.6893};

static const std::vector<std::vector<double>> expect_derivative{
{-25.0762, -0, -0, -0, -0, -0, -0, -0, -0, -0}, 
{-0, -8.7775, -0, -0, -0, -0, -0, -0, -0, -0}, 
{-0, -0, -6.36414, -0, -0, -0, -0, -0, -0, -0}, 
{-0, -0, -0, -0.904231, -0, -0, -0, -0, -0, -0}};

void CrossEntropyLossTest::test()
{
    CrossEntropyLoss cross_entropy_loss{};
    double loss{};
    cross_entropy_loss.forward(input, labels, loss);
    // LogInfo() << "loss: " << loss;
    auto derivative = cross_entropy_loss.backward();
    // LogInfo() << "derivative: " << derivative;

    assert(Tensor::equal(loss, expect_loss, 1e-5));
    assert(Tensor::equal(derivative, expect_derivative, 1e-4));
    LogInfo() << "CrossEntropyLossTest test success.";
}

void CrossEntropyLossTest::performance_test()
{
    CrossEntropyLoss cross_entropy_loss{};
    auto input = Tensor::reshape(std::vector<double>{0.1, 0.2, 0.3, 0.4, 0.0}, 1, 5);
    auto expect = Tensor::reshape(std::vector<double>{1, 0, 0, 0, 0}, 1, 5);
    double loss{};

    int64_t epochs{1000000};
    MEASURE_CPU_TIME("CrossEntropyLossTest", epochs, {
        cross_entropy_loss.forward(input, expect, loss);
        cross_entropy_loss.backward();
    });
}